import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
import torchvision.transforms as transforms
import torchvision
import torch.utils.data.dataloader
from ResNet18 import ResNet18
import torch

# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# tensorBoard
writer = SummaryWriter("logs")

# 超参数设置
epoches = 30
batch_size = 128
learning_rate = 0.01

# 准备数据集预处理
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

train_dataset = torchvision.datasets.CIFAR10(root='./resources', train=True, download=True, transform=transform_train)

test_dataset = torchvision.datasets.CIFAR10(root='./resources', train=False, download=True, transform=transform_test)

# num_works 多少个工作进程，可是找batch的速度变快
train_loader = torch.utils.data.dataloader.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)

test_loader = torch.utils.data.dataloader.DataLoader(test_dataset, batch_size=100, shuffle=False, num_workers=2)

# 标签值
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# 模型定义
net = ResNet18().to(device)

# 定义损失函数和优化方式
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.0001)
# 学习率调整，在第10个epoch中改变为0.01,20个epoch为0.001
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [10, 20], gamma=0.1)

# 训练
if __name__ == "__main__":
    train_lossList = []
    test_lossList = []
    for epoch in range(epoches):
        print('\nEpoch:%d' % (epoch+1))
        # 进入训练模式
        net.train()
        correct = 0.0
        total = 0.0
        for i, (inputs, labels) in enumerate(train_loader):
            length = len(train_loader)
            inputs = inputs.to(device)
            labels = labels.to(device)

            optimizer.zero_grad()

            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (labels == predicted).sum().item()

            writer.add_scalar("loss", loss, i + 1 + epoch * length)
            print('[epoch:{}, iter:{}] Loss: {:.3f}|Acc: {:.3f}%'
                  .format(epoch + 1, (i + 1 + epoch * length), loss.item(), 100 * (correct / total)))

        # 测试准确率
        print("Waiting Test!")
        with torch.no_grad():
            test_correct = 0
            test_total = 0
            for (images, labels) in test_loader:
                images = images.to(device)
                labels = labels.to(device)
                outputs = net(images)
                _, predicted = torch.max(outputs.data, 1)
                test_total += labels.size(0)
                test_correct += (predicted == labels).sum()
            test_accuracy = 100*test_correct/test_total
            writer.add_scalar("test_accuracy", test_accuracy, epoch)
            print('测试分类准确率为{:.3f}'.format(test_accuracy))

        scheduler.step()
    print('训练完成！')
    writer.close()


